#!/usr/bin/env python
import numpy as np
import XmlTools
from maputility2 import maptools, skyutils as maputils

def getPositions(infile, min_ts=20., min_angsep=3., min_radius=1., max_radius=10., ts_weight=True):
    if min_radius>=max_radius:
        raise IOError('Invalid values: min_radius=%s, max_radius=%s!' % (min_radius, max_radius))

    m = maptools.mapOpen(infile)
    if m._wcs.coord is None:
        raise RuntimeError('Mapfile %s does not have coordinate system!' % infile)

    mask0 = (m.map >= min_ts) # ts Limit, shape in (nypix, nxpix)
    skyctr = np.array(m.pix2sky((m.nxpix-1.)/2., (m.nypix-1.)/2.))
    angsep_center = maputils.angsep(skyctr, m.skies._.T).reshape(m.nxpix, m.nypix).T
    mask1 = (angsep_center>=min_radius)&(angsep_center<max_radius) # radius Limit, shape in (nypix, nxpix)
    mask = mask0 & mask1

    if mask.any():
        skymaps = m.skies._.reshape(2, m.nxpix, m.nypix).T # (nypix, nxpix, 2)
        xys = skymaps[mask, :].reshape(-1, 2) # (npix, 2)

        angseps = np.triu(np.atleast_2d(maputils.angsep(xys, xys)), k=1)
        mask2 = (angseps<=min_angsep)&(angseps>0.) # remove too close points
        if mask2.any():
            nonzeros = np.nonzero(mask2)
            vals = m.map[mask].ravel()
            maskxys = np.full_like(xys, dtype=bool, fill_value=True)
            for i, j in zip(*nonzeros):
                val1, val2 = vals[i], vals[j]
                if val1 <= val2:
                    maskxys[i, :] = False
                else:
                    maskxys[j, :] = False
            xys = xys[maskxys].reshape(-1, 2)

        px, py = m.sky2pix(xys[:, 0], xys[:, 1]) # npix
        #print(px.min(), px.max(), py.min(), py.max())
        ts_xys = m.map[py, px].reshape(-1,1)
        #print(px, py, ts_xys)

        if ts_weight: # weight the position with ts^2
            pxl = np.maximum(px-1, 0)
            if (pxl<0.).any(): raise RuntimeError('pix_left<0')
            left = np.atleast_2d(m.pix2sky(pxl, py, is_squeeze=False)).T
            ts_left = m.map[py, pxl].reshape(-1, 1)

            pxr = np.minimum(px+1, m.nxpix-1)
            if (pxr>=m.nxpix).any(): raise RuntimeError('pix_right>=nxpix')
            right = np.atleast_2d(m.pix2sky(pxr, py, is_squeeze=False)).T
            ts_right = m.map[py, pxr].reshape(-1, 1)

            pyd = np.maximum(py-1, 0)
            if (pyd<0.).any(): raise RuntimeError('pix_down<0')
            down = np.atleast_2d(m.pix2sky(px, pyd, is_squeeze=False)).T
            ts_down = m.map[pyd, px].reshape(-1, 1)

            pyu = np.minimum(py+1, m.nypix-1)
            if (pyu>=m.nypix).any(): raise RuntimeError('pix_up>=nypix')
            up = np.atleast_2d(m.pix2sky(px, pyu, is_squeeze=False)).T
            ts_up = m.map[pyu, px].reshape(-1, 1)

            #print(xys.shape, ts_xys.shape, left.shape, ts_left.shape)

            xys_new = (xys*ts_xys**2 + left*ts_left**2 + right*ts_right**2 + up*ts_up**2 + down*ts_down**2) /\
                      (ts_xys**2 + ts_left**2 + ts_right**2 + ts_up**2 + ts_down**2)
            xys = xys_new

        if m._wcs.coord == 'GAL':
            xys = np.atleast_2d(maputils.GAL2CEL(xys[:, 0], xys[:, 1])).T
        return xys, ts_xys
    else:
        return None


def addPts(pos, ts, infile=None, outfile=None, newsrc_prefix='newpt', min_index=0, prefactor_scale=1e-13, makeReg=True):
    if infile:
        outfile = outfile or '.'.join(infile.split('.')[:-1])+'_addNewPts.xml'
        x = XmlTools.SMKernel.SourceModel(infile, raiseErr=True)
    else:
        outfile = outfile or 'newpts.xml'
        x = XmlTools.SMKernel.SourceModel(raiseErr=False)

    pos_array = np.atleast_2d(pos)
    ts_array = np.array(ts).ravel()
    #print(pos_array, ts_array)
    for isrc in range(pos_array.shape[0]):
        src_name = '%s_%s'%(newsrc_prefix, min_index+isrc)
        src_pos = pos_array[isrc, :]
        print('Add a new point %s at FK5(%s, %s)' % (src_name, src_pos[0], src_pos[1]))
        nm = XmlTools.Models.PowerLaw(src_name, ra=src_pos[0], dec=src_pos[1])
        nm.spectrum.Prefactor.scale = prefactor_scale
        nm['TS_value'] = ts_array[isrc]
        x[None] = nm

    x.writeTo(outfile)
    if makeReg:
        x.makeRegion('ROI_'+newsrc_prefix+'.reg')

def cli():
    import argparse

    parser = argparse.ArgumentParser(description='Add new pts according to the tsmap.')
    parser.add_argument("tsmap", type=str, help='Input TS file')
    parser.add_argument("-srcmdl", type=str, default='', help='Input source model file')
    parser.add_argument("-outfile", type=str, default='', help='Output source model file')
    parser.add_argument("-min_radius", type=float, default=1., help='Minimum radius[deg] to add new pts')
    parser.add_argument("-max_radius", type=float, default=10., help='Maximum radius[deg] to add new pts')
    parser.add_argument("-min_ts", type=float, default=20., help='Minimum TS value to add new point sources')
    parser.add_argument("-min_as", type=float, default=2., help='Minimum angular separation between two added point sources')
    parser.add_argument("-src_prefix", type=str, default='newpt', help='The prefix of new point sources')
    parser.add_argument("-min_index", type=int, default=0, help='From what number the name of new src is given')
    parser.add_argument("-makeRegion", action='store_false', default=True, help='Whether to make region file (Default: True)')
    parser.add_argument("-ts_weight", action='store_false', default=True, help='Whether to weight position with ts_value (Default: True)')
    args = parser.parse_args()

    out = getPositions(args.tsmap,
                       min_ts=args.min_ts, min_angsep=args.min_as, min_radius=args.min_radius, max_radius=args.max_radius,
                       ts_weight=args.ts_weight)
    if out is None:
        print('No new point sources are added!')
    else:
        xys, ts = out
        addPts(xys, ts, infile=args.srcmdl, outfile=args.outfile,
               newsrc_prefix=args.src_prefix, min_index=args.min_index,
               makeReg=args.makeRegion)

if __name__ == '__main__': cli()
